"""
PRISM  - 


"""

import sys
import os
sys.path.insert(0, os.path.dirname(os.path.dirname(os.path.abspath(__file__))))

from src.optimization.hyperparameter_search import HyperparameterOptimizer
import logging

logging.basicConfig(level=logging.INFO)
logger = logging.getLogger(__name__)


def run_quick_optimization():
    """
    

    
    1. Stage 1 5
    2. Stage 2 2-4
    3. 10
    """

    print("=" * 80)
    print("PRISM  - ")
    print("=" * 80)
    print("\n")
    print("1. Stage 1 proposals")
    print("2. Stage 2 10 trials")
    print("3. ")
    print("\n2-4")
    print("=" * 80)

    response = input("\n(y/n): ")
    if response.lower() != 'y':
        print("")
        return

    # ===== 1: Stage 1  =====
    print("\n\n" + "=" * 80)
    print(" 1/3: Stage 1 Proposals ")
    print("=" * 80)

    optimizer_s1 = HyperparameterOptimizer(optimization_phase='stage1')
    results_s1 = optimizer_s1.optimize()

    print("\n Stage 1 ")
    print(f": {results_s1['best_config']['threshold']}")

    # 
    response = input("\nStage 2(y/n): ")
    if response.lower() != 'y':
        print("")
        return

    # ===== 2: Stage 2  =====
    print("\n\n" + "=" * 80)
    print(" 2/3: Stage 2 ")
    print("=" * 80)
    print("2-4")

    optimizer_s2 = HyperparameterOptimizer(optimization_phase='stage2')
    results_s2 = optimizer_s2.optimize(n_trials=10)  # 10 trials

    print("\n Stage 2 ")
    print(f" mAP: {results_s2['best_value']:.4f}")

    # 
    print("\n...")
    optimizer_s2.visualize_results()

    # 
    response = input("\n(y/n): ")
    if response.lower() != 'y':
        print("")
        return

    # ===== 3:  =====
    print("\n\n" + "=" * 80)
    print(" 3/3: ")
    print("=" * 80)

    optimizer_deploy = HyperparameterOptimizer(optimization_phase='deployment')
    results_deploy = optimizer_deploy.optimize()

    print("\n ")
    print(f"Stage 1 : {results_deploy['best_config']['stage1_threshold']}")
    print(f"Stage 2 : {results_deploy['best_config']['stage2_threshold']}")
    print(f" mAP: {results_deploy['best_config']['map']:.4f}")

    # =====  =====
    print("\n\n" + "=" * 80)
    print(" ")
    print("=" * 80)

    print("\n ")
    print(f"\nStage 1:")
    print(f"  Proposals: {results_s1['best_config']['threshold']}")

    print(f"\nStage 2:")
    for key, value in results_s2['best_params'].items():
        print(f"  {key}: {value}")

    print(f"\n:")
    print(f"  Stage 1: {results_deploy['best_config']['stage1_threshold']}")
    print(f"  Stage 2: {results_deploy['best_config']['stage2_threshold']}")

    print(f"\n:")
    print(f"  mAP: {results_deploy['best_config']['map']:.4f}")
    print(f"  Precision: {results_deploy['best_config']['precision']:.3f}")
    print(f"  Recall: {results_deploy['best_config']['recall']:.3f}")

    print("\n : optimization_results/")
    print("=" * 80)


def run_custom_optimization():
    """"""

    print("=" * 80)
    print("PRISM  - ")
    print("=" * 80)

    print("\n:")
    print("1. Stage 1 ")
    print("2. Stage 2 ")
    print("3. ")
    print("4. ")

    choice = input("\n (1-4): ")

    phase_map = {
        '1': 'stage1',
        '2': 'stage2',
        '3': 'deployment',
        '4': 'end_to_end'
    }

    if choice not in phase_map:
        print("")
        return

    phase = phase_map[choice]

    # trials
    n_trials = 20
    if phase in ['stage2', 'end_to_end']:
        n_trials_input = input(f"\ntrials20: ")
        if n_trials_input:
            n_trials = int(n_trials_input)

    # 
    optimizer = HyperparameterOptimizer(optimization_phase=phase)

    if phase in ['stage2', 'end_to_end']:
        results = optimizer.optimize(n_trials=n_trials)

        # 
        print("\n...")
        optimizer.visualize_results()
    else:
        results = optimizer.optimize()

    print("\n ")
    print(f": optimization_results/")


if __name__ == '__main__':
    import argparse

    parser = argparse.ArgumentParser(description='PRISM ')
    parser.add_argument('--mode', type=str, default='quick',
                       choices=['quick', 'custom'],
                       help='quick custom')

    args = parser.parse_args()

    if args.mode == 'quick':
        run_quick_optimization()
    else:
        run_custom_optimization()
